from api.extension import return_code
from django.core.cache import cache as default_cache
from rest_framework import status
from rest_framework.exceptions import APIException
from rest_framework.throttling import SimpleRateThrottle


class ThrottledException(APIException):
    status_code = status.HTTP_429_TOO_MANY_REQUESTS
    default_code = 'throttled'


class NewsCreateRateThrottle(SimpleRateThrottle):
    cache = default_cache
    cache_format = 'throttle_%(scope)s_%(ident)s'
    scope = 'user'  # 构造缓存中的key
    # 设置访问频率
    THROTTLE_RATES = {'user': '1/5m'}  # s m h d 重写方法只支持单字母

    def parse_rate(self, rate):
        if rate is None:
            return None, None
        num, period = rate.split('/')
        num_requests = int(num)
        count = int(period[:-1])
        duration = {'s': 1, 'm': 60, 'h': 3600, 'd': 86400}[period[-1]]
        return num_requests, duration * count

    def get_cache_key(self, request, view):
        ident = request.user.pk
        return self.cache_format % {'scope': self.scope, 'ident': ident}

    def throttle_failure(self):
        wait = self.wait()
        detail = {
            'code': return_code.TOO_MANY_REQUESTS,
            'data': '访问频率限制',
            'detail': f'需等待{int(wait)}秒后才能访问'
        }
        raise ThrottledException(detail)

    def throttle_success(self):
        return True

    def done(self):
        """ 数据库执行成功后调用，本来是限流放行后即执行 """
        self.history.insert(0, self.now)
        self.cache.set(self.key, self.history, self.duration)
